
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import readline from "readline/promises";
import { OpenAI } from "openai";
import { ChatCompletionMessageParam, ChatCompletionMessageToolCall, ChatCompletionTool } from "openai/resources.mjs";
import dotenv from "dotenv";

dotenv.config();

const DEEPSEEK_API_KEY = process.env.DEEPSEEK_API_KEY;
if (!DEEPSEEK_API_KEY) {
  throw new Error("DEEPSEEK_API_KEY is not set");
}

class MCPClient {
  private mcp: Client;
  private openai: OpenAI;
  private transport: StdioClientTransport | null = null;
  private tools: ChatCompletionTool[] = [];

  constructor() {
    this.openai = new OpenAI({
      baseURL: "https://api.deepseek.com/v1",
      apiKey: DEEPSEEK_API_KEY,
    });
    this.mcp = new Client({ name: "swagger2-doc-mcp", version: "1.0.0" });
  }

  
  async connectToServer(serverScriptPath: string) {
    try {
      const isJs = serverScriptPath.endsWith(".js");
      const isPy = serverScriptPath.endsWith(".py");
      if (!isJs && !isPy) {
        throw new Error("Server script must be a .js or .py file");
      }
      const command = isPy
        ? process.platform === "win32"
          ? "python"
          : "python3"
        : process.execPath
      
      this.transport = new StdioClientTransport({
        command,
        args: [serverScriptPath],
      });
      this.mcp.connect(this.transport);
      
      const toolsResult = await this.mcp.listTools();
      this.tools = toolsResult.tools.map((tool) => {
        return {
          type: "function",
          function: {
            name: tool.name,
            description: tool.description,
            parameters: tool.inputSchema,
          }
        };
      }) as ChatCompletionTool[];
      console.log(
        "Connected to server with tools:",
        this.tools.map((tool) => tool.function.name)
      );
    } catch (e) {
      console.log("Failed to connect to MCP server: ", e);
      throw e;
    }
  }

  async processQuery(query: string) {
    const messages: ChatCompletionMessageParam[] = [
      {
        role: "user",
        content: query,
      },
    ];
  
    const response = await this.openai.chat.completions.create({
      model: "deepseek-chat",
      max_tokens: 1000,
      messages,
      tools: this.tools,
    });
  
    const finalText: string[] = [];
    const toolResults: any[] = [];
  
    for (const content of response.choices) {
      if (content.finish_reason === "stop") {
        finalText.push(content.message.content ?? "");
      } else if (content.finish_reason === "tool_calls" && content.message.tool_calls?.length && content.message.tool_calls?.length > 0) {
        const tool = content.message.tool_calls[0] as ChatCompletionMessageToolCall;
        const toolName = tool.function.name
        const toolArgs = tool.function.arguments
        const result = await this.mcp.callTool({
          name: toolName,
          arguments: JSON.parse(toolArgs),
        });
        toolResults.push(result);
        finalText.push(
          `[Calling tool ${toolName} with args ${JSON.stringify(toolArgs)}]`
        );
  
        messages.push({
          role: "user",
          content: result.content as string,
        });
  
        const response = await this.openai.chat.completions.create({
          model: "deepseek-chat",
          max_tokens: 1000,
          messages
        });
  
        finalText.push(
          response.choices[0].finish_reason === "stop" ? response.choices[0].message.content ?? "" : ""
        );
      }
    }
  
    return finalText.join("\n");
  }

  async chatLoop() {
    const rl = readline.createInterface({
      input: process.stdin,
      output: process.stdout,
    });
  
    try {
      console.log("\nMCP Client Started!");
      console.log("Type your queries or 'quit' to exit.");
  
      while (true) {
        const message = await rl.question("\nQuery: ");
        if (message.toLowerCase() === "quit") {
          break;
        }
        const response = await this.processQuery(message);
        console.log("\n" + response);
      }
    } finally {
      rl.close();
    }
  }
  
  async cleanup() {
    await this.mcp.close();
  }
}


async function main() {
  if (process.argv.length < 3) {
    console.log("Usage: node index.ts <path_to_server_script>");
    return;
  }
  const mcpClient = new MCPClient();
  try {
    await mcpClient.connectToServer(process.argv[2]);
    await mcpClient.chatLoop();
  } catch(error) {
    console.log(error)
  } finally {
    await mcpClient.cleanup();
    process.exit(0);
  }
}

main();